# 导入所需模块
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    Settings,
    PromptTemplate,
)
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding

# 1. 设置 Ollama 模型（LLM 和 Embedding）
Settings.llm = Ollama(
    model="qwen3:8b",
    request_timeout=300.0,
    base_url="http://localhost:11434"
)

Settings.embed_model = OllamaEmbedding(
    model_name="nomic-embed-text:v1.5",
    base_url="http://localhost:11434",
    ollama_additional_kwargs={"keep_alive": 1}
)

# 2. 加载文档
documents = SimpleDirectoryReader("data").load_data()

# 3. 构建索引
index = VectorStoreIndex.from_documents(documents)

# 4. 定义你自己的 Prompt 模板
CUSTOM_PROMPT_TEMPLATE = """
你是一个专业的智能助手，请根据以下提供的上下文信息回答用户的问题。
请遵循以下规则：
1. 只基于上下文回答，不要编造信息。
2. 如果上下文中没有答案，请回答：“抱歉，根据提供的资料无法回答该问题。”
3. 回答尽量简洁清晰，使用中文。

上下文信息如下：
---------------------
{{context_str}}
---------------------

用户问题：
{{query_str}}

请开始回答：

/nothink
"""

# 创建 PromptTemplate 实例
custom_prompt = PromptTemplate(CUSTOM_PROMPT_TEMPLATE)

# 5. 创建查询引擎，并传入自定义模板
query_engine = index.as_query_engine(
    streaming=True,
    text_qa_template=custom_prompt,   # 第一次回答用的模板
    refine_template=custom_prompt,    # 后续 refine 用的模板（当上下文太多需分段处理）
    # use_async=True  # 可选：异步处理提升性能
)

# 6. 开始问答
while True:
    question = input("\n问一个问题 (输入 'quit' 退出): ")
    if question.lower() == "quit":
        break
    response = query_engine.query(question)
    print("\n回答: ")
    response.print_response_stream()